from langchain.chat_models import ChatOpenAI
from langchain.agents import AgentType, initialize_agent, Tool
from langchain.memory import ConversationBufferMemory
import requests
from dotenv import load_dotenv
import os

# 加载环境变量（包含接口地址和密钥）
load_dotenv()

# 1. 天气查询工具函数（核心工具）
def get_weather(city: str) -> str:
    """查询指定城市的实时天气，输入为城市名称（如北京、上海）"""
    # 调用OpenWeatherMap API
    weather_url = f"http://api.openweathermap.org/data/2.5/weather?q={city}&appid={os.getenv('WEATHER_API_KEY')}&units=metric"
    try:
        response = requests.get(weather_url, timeout=10)
        data = response.json()
        # 处理查询失败的情况
        if data.get("cod") != 200:
            return f"天气查询失败：{data.get('message', '未知错误')}"
        # 提取天气信息（温度和描述）
        temp = data["main"]["temp"]
        desc = data["weather"][0]["description"]
        return f"{city} 当前气温：{temp}°C，天气：{desc}"
    except Exception as e:
        return f"查询过程出错：{str(e)}"

# 2. 定义工具列表（告诉Agent可用的工具及用途）
tools = [
    Tool(
        name="WeatherQuery",  # 工具名称（Agent调用时会引用）
        func=get_weather,     # 绑定工具函数
        description="用于查询任意城市的实时天气，当用户询问天气相关问题时，必须调用该工具获取数据，不能凭经验回答"
    )
]

# 3. 初始化本地模型（适配OpenAI兼容接口）
# 仅保留接口支持的参数（model、max_tokens、stream）
local_llm = ChatOpenAI(
    model="DeepSeekR1",  # 必须与接口要求的model名称一致
    max_tokens=500,      # 接口支持的参数，控制最大输出长度
    # 注意：接口未提及temperature等参数，如需添加需先确认接口是否支持，否则删除
    # temperature=0.5  # 若接口不支持，保留会导致错误，暂时移除
)

# 4. 配置对话记忆（让Agent记住历史对话）
memory = ConversationBufferMemory(
    memory_key="chat_history",  # 记忆在对话中的存储键名
    return_messages=True        # 返回消息对象，适配对话型Agent
)

# 5. 初始化Agent（整合模型、工具、记忆）
agent = initialize_agent(
    tools=tools,
    llm=local_llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,  # 支持记忆的对话型Agent
    verbose=True,  # 打印思考过程（调试用，上线可设为False）
    memory=memory
)

# 6. 测试Agent功能
if __name__ == "__main__":
    print("===== 测试1：查询天气 =====")
    agent.run("查询Singapore的天气")
    
    print("\n===== 测试2：基于历史对话提问 =====")
    agent.run("刚才查询的城市天气如何？适合出门吗？")
